import rclpy
from rclpy.node import Node
from geometry_msgs.msg import (
    PoseWithCovarianceStamped,
    PoseStamped,
    PointStamped,
    Vector3,
)
from nav_msgs.msg import Path
from visualization_msgs.msg import Marker
from std_msgs.msg import ColorRGBA
import numpy as np

import numpy as np


def bezier_curve(p0, p1, p2, p3, num_points=20):
    """
    计算四阶贝塞尔曲线上的点

    :param p0: 控制点0 (x0, y0)
    :param p1: 控制点1 (x1, y1)
    :param p2: 控制点2 (x2, y2)
    :param p3: 控制点3 (x3, y3)
    :param num_points: 曲线上的点的数量
    :return: 二维数组，包含曲线上的点
    """
    t = np.linspace(0, 1, num_points)
    curve = (
        (1 - t) ** 3 * np.array(p0)
        + 3 * (1 - t) ** 2 * t * np.array(p1)
        + 3 * (1 - t) * t**2 * np.array(p2)
        + t**3 * np.array(p3)
    )
    return [(float(point[0]), float(point[1])) for point in curve]


class DemoNode(Node):
    def __init__(self):
        super().__init__("demo_node")
        self.get_logger().info("Demo node started")
        self.init_pose_subscriber = self.create_subscription(
            PoseWithCovarianceStamped, "initialpose", self.init_pose_callback, 10
        )
        self.target_pose_subscriber = self.create_subscription(
            PoseStamped, "goal_pose", self.target_pose_callback, 10
        )
        self.mid_point_subscriber = self.create_subscription(
            PointStamped, "clicked_point", self.mid_point_callback, 10
        )
        self.start_marker_publisher = self.create_publisher(Marker, "start_marker", 10)
        self.target_marker_publisher = self.create_publisher(
            Marker, "target_marker", 10
        )
        self.cp_publisher = self.create_publisher(Marker, "cp", 10)
        self.path_publisher = self.create_publisher(Path, "path", 10)
        self.hasStart = False
        self.hasTarget = False

    def init_pose_callback(self, msg: PoseWithCovarianceStamped):
        self.get_logger().info("Initial pose received")
        self.start_x = msg.pose.pose.position.x
        self.start_y = msg.pose.pose.position.y
        self.start_angle = msg.pose.pose.orientation.z
        self.get_logger().info(
            "x:{}, y:{}, angle:{}".format(self.start_x, self.start_y, self.start_angle)
        )
        self.start_marker_publisher.publish(
            Marker(
                header=msg.header,
                ns="start_marker",
                id=0,
                type=Marker.ARROW,
                action=Marker.ADD,
                pose=msg.pose.pose,
                scale=Vector3(x=5.0, y=1.0, z=1.0),
                color=ColorRGBA(r=0.0, g=0.0, b=1.0, a=1.0),
            )
        )
        self.hasStart = True
        if self.hasStart and self.hasTarget:
            self.create_path()

    def target_pose_callback(self, msg: PoseStamped):
        self.get_logger().info("Target pose received")
        self.target_x = msg.pose.position.x
        self.target_y = msg.pose.position.y
        self.target_angle = msg.pose.orientation.z
        self.get_logger().info(
            "x:{}, y:{}, angle:{}".format(
                self.target_x, self.target_y, self.target_angle
            )
        )
        self.target_marker_publisher.publish(
            Marker(
                header=msg.header,
                ns="start_marker",
                id=1,
                type=Marker.ARROW,
                action=Marker.ADD,
                pose=msg.pose,
                scale=Vector3(x=5.0, y=1.0, z=1.0),
                color=ColorRGBA(r=0.0, g=1.0, b=0.0, a=1.0),
            )
        )
        self.hasTarget = True
        if self.hasStart and self.hasTarget:
            self.create_path()

    def mid_point_callback(self, msg: PointStamped):
        self.get_logger().info("Mid point received")
    def create_path(self):
        init_vec_ctrl = 2
        target_vec_ctrl = 10
        p0 = (self.start_x, self.start_y)
        p1 = (
            self.start_x + init_vec_ctrl * np.sin(self.start_angle),
            self.start_y + init_vec_ctrl * np.cos(self.start_angle),
        )
        p2 = (
            self.target_x - target_vec_ctrl * np.sin(self.target_angle),
            self.target_y - target_vec_ctrl * np.cos(self.target_angle),
        )
        p3 = (self.target_x, self.target_y)
        points = bezier_curve(p0, p1, p2, p3, num_points=20)
        path = Path()
        path.header.frame_id = "map"
        path.header.stamp = self.get_clock().now().to_msg()
        for point in points:
            pose = PoseStamped()
            pose.header.frame_id = "map"
            pose.header.stamp = self.get_clock().now().to_msg()
            pose.pose.position.x = point[0]
            pose.pose.position.y = point[1]
            pose.pose.position.z = 0.0
            # pose.pose.orientation.z = point[2]
            path.poses.append(pose)
        self.get_logger().info("Path created")
        self.path_publisher.publish(path)


def main(args=None):
    rclpy.init(args=args)
    demo_node = DemoNode()
    rclpy.spin(demo_node)
    demo_node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()
